import json
import os

from dotenv import load_dotenv
from influxdb_client import InfluxDBClient, Point, WritePrecision
from influxdb_client.client.write_api import SYNCHRONOUS


# 配置你的InfluxDB服务器信息
BATCH_SIZE = 10
WRITE_PRECISION = WritePrecision.NS


def create_point(kline):
    timestamp_ns = kline["k"]["t"] * 1_000_000  # Convert to nanoseconds
    print(f"Timestamp (ns): {timestamp_ns}")

    point = (
        Point("kline_data")
        .tag("symbol", kline["s"])
        .tag("interval", kline["k"]["i"])
        .field("o", float(kline["k"]["o"]))
        .field("c", float(kline["k"]["c"]))
        .field("h", float(kline["k"]["h"]))
        .field("l", float(kline["k"]["l"]))
        .field("v", float(kline["k"]["v"]))
        .time(timestamp_ns, WritePrecision.NS)
    )
    return point


class InfluxDBManager:
    """InfluxDB连接管理类"""

    def __init__(self):
        load_dotenv()
        self.INFLUXDB_URL = os.getenv('INFLUXDB_URL')
        self.ORG = os.getenv('INFLUXDB_ORG')
        self.BUCKET_NAME = os.getenv('INFLUXDB_BUCKET_NAME')
        self.TOKEN = os.getenv('INFLUXDB_TOKEN')
        self.client = None

    def get_client(self):
        """获取InfluxDB客户端实例，如果尚未创建则初始化"""
        if not self.client:
            self.client = InfluxDBClient(url=self.INFLUXDB_URL, token=self.TOKEN, org=self.ORG)
        return self.client

    def write_points(self, points):
        """写入数据点到InfluxDB"""
        try:
            if not points:
                print("No points to write.")
                return

            print(f"Writing {points[0]} to InfluxDB...")

            write_api = self.get_client().write_api(write_options=SYNCHRONOUS)

            # 调试信息：打印所有点的数据
            for point in points:
                print(f"Point: {point.to_line_protocol()}")

            # 写入数据
            write_api.write(bucket=self.BUCKET_NAME, record=points)

            print(f"Wrote {len(points)} points to InfluxDB.")
        except Exception as e:
            print(f"Failed to write points to InfluxDB: {e}")

    def close(self):
        """关闭数据库连接"""
        if self.client:
            self.client.close()
            self.client = None


# 使用示例
def write_klines_to_influxdb(klines, manager):
    """
    使用InfluxDBManager将K线数据写入InfluxDB。

    :param klines: 包含K线数据的列表。
    :param manager: InfluxDBManager实例。
    """
    points = []
    for kline in klines:
        print(f"Processing kline: {json.dumps(kline)}")
        point = create_point(kline)
        points.append(point)
        if len(points) >= BATCH_SIZE:
            manager.write_points(points)
            points = []
    if points:
        manager.write_points(points)
